package be.android.forap.crypto;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.SecretKeySpec;


public class AesOcb {

	static String algorithm = "AES";
	public static final int BLOCK_BITS  = 128;

	protected static final int  PRE_COMP_BLOCKS = 32;
	public static final int BLOCK_SIZE  = (BLOCK_BITS >>> 3);
	protected static final int  TAG_LENGTH = BLOCK_SIZE;

	protected byte[][]  L       = new byte[PRE_COMP_BLOCKS][BLOCK_SIZE]; // Precomputed L(i) values.
	protected byte[]    L_inv   = new byte[BLOCK_SIZE]; // Precomputed L/x value
	protected byte[]    tmp     = new byte[BLOCK_SIZE];
	protected byte[]    offset  = new byte[BLOCK_SIZE]; // Offset (Z[i]) for current block
	protected byte[]    chksum  = new byte[TAG_LENGTH]; // Checksum for computing tag
	
	/***
	 *  function ocb-aes-encrypt(K, M, Nonce)
    begin
        Offset = AES(K, Nonce xor L) 
        Checksum = 0
        for i = 1 to m-1 do begin
            Offset = Offset xor L(ntz(i))
            Checksum = Checksum xor M[i]
            C[i] = Offset xor AES(K, M[i] xor Offset) 
        end
        Offset = Offset xor L(ntz(m))
        Pad = AES (K, len(M[m]) xor L(-1) xor Offset)
        C[m] = M[m] xor (the first |M[m]| bits of Pad)
        Checksum = Checksum xor Pad xor C[m]0*
        FullTag = AES(K, Checksum xor Offset) 
        Tag = a prefix of FullTag (of the desired length)
        return C[1]... C[m-1] C[m] Tag 
    end
	 */

	public byte[] encrypt(byte[] pt, int ptLen, byte[] nonce, byte[] tag){

		int ptPos = 0;

		// Create ciphertext
		byte[] ct = new byte[ptLen];
		int ctPos = 0;

		Arrays.fill(chksum, (byte)0);                       // Zero the checksum
		xorBlock(offset, L[0], nonce);          			// Calculate R, aka Z[0]
		offset = encryptAes(offset);


		// Process blocks 1 .. m-1
		int i;
		for (i = 1; ptLen > BLOCK_SIZE; i++) {
			xorBlock(offset, offset, L[ntz(i)]);			// Update the offset (Z[i] from Z[i-1])
			xorBlock(chksum, chksum, pt, ptPos);			// Update the checksum
			xorBlock(tmp, offset, pt, ptPos);				// xor the plaintext block with Z[i]
			tmp = encryptAes(tmp);							// Encipher the block
			xorBlock(ct, ctPos, offset, tmp);				// xor Z[i] again, writing result to ciphertext pointer
			ptLen -= BLOCK_SIZE;
			ptPos += BLOCK_SIZE;
			ctPos += BLOCK_SIZE;
		}

		// Process block m
		xorBlock(offset, offset, L[ntz(i)]);                // Update the offset (Z[m] from Z[m-1])
		xorBlock(tmp, offset, L_inv);                       // xor L . x^{-1} and Z[m]
		tmp[BLOCK_SIZE - 1] ^= (byte)(ptLen << 3);          // Add in final block bit-length
		tmp = encryptAes(tmp);

		for (int t = 0; t < ptLen; t++) {
			ct[ctPos + t] = (byte)(pt[ptPos + t] ^ tmp[t]); // xor pt with block-cipher output to get ct
			tmp[t] = pt[ptPos + t];                         // Add to checksum the ptLen bytes of plaintext...
		}
		xorBlock(chksum, chksum, tmp);                      // ... followed by the last (16 - ptLen) bytes of block-cipher output

		// Calculate tag
		xorBlock(chksum, chksum, offset);
		tmp = encryptAes(chksum);
		System.arraycopy(tmp, 0, tag, 0, TAG_LENGTH);

		return ct;
	}

	/**
	 * Given a nonce starting at offset noncePos of the nonce byte array,
	 * decrypt ctLen elements of byte array ct starting at offset ctPos, and
	 * verifying a tag starting at offset tagPos of the tag byte array.
	 *
	 * @param       ct          ciphertext buffer.
	 * @param       ctPos       start index of ciphertext on ct.
	 * @param       ctLen       byte length of ciphertext.
	 * @param       nonce       nonce array (BLOCK_SIZE bytes).
	 * @param       noncePos    start index of nonce.
	 * @param       tag         input tag buffer.
	 * @param       tagPos      start index of tag.
	 *
	 * @return      the resulting plaintext, a byte array of same length as ct
	 *              if decryption is successfull, or else null if the tag does
	 *              not correctly verify.
	 * 
	 * @warning     ct, nonce, and tag must not overlap.
	 */
	public final byte[] decrypt(byte[] ct, int ctLen, byte[] nonce,  byte[] tag)
			throws IllegalArgumentException {

		int ctPos = 0;
		// Create plaintext
		byte[] pt = new byte[ctLen];
		int ptPos = 0;

		Arrays.fill(chksum, (byte)0);               // Zero checksum
		xorBlock(offset, L[0], nonce);    			// Calculate R, aka Z[0]
		offset = encryptAes(offset);

		// Process blocks 1 .. m-1
		int i;
		for (i = 1; ctLen > BLOCK_SIZE; i++) {
			xorBlock(offset, offset, L[ntz(i)]);            // Update the offset (Z[i] from Z[i-1])
			xorBlock(tmp, offset, ct, ctPos);               // xor ciphertext block with Z[i]
			tmp = decryptAes(tmp);                          // Decipher the next block-cipher block
			
			xorBlock(pt, ptPos, offset, tmp);               // xor Z[i] again, writing result to plaintext pointer
			xorBlock(chksum, chksum, pt, ptPos);            // Update the checksum
			ctLen -= BLOCK_SIZE;
			ctPos += BLOCK_SIZE;
			ptPos += BLOCK_SIZE;
		}

		// Process block m
		xorBlock(offset, offset, L[ntz(i)]);                // Update the offset (Z[m] from Z[m-1])
		xorBlock(tmp, offset, L_inv);                       // xor L . x^{-1} and Z[m]
		tmp[BLOCK_SIZE - 1] ^= (byte)(ctLen << 3);          // Add in final block bit-length
		tmp = encryptAes(tmp);

		for (int t = 0; t < ctLen; t++) {
			pt[ptPos + t] = (byte)(ct[ctPos + t] ^ tmp[t]); // xor ct with block-cipher output to get pt
			tmp[t] = pt[ptPos + t];                         // Add to checksum the ctLen bytes of plaintext...
		}
		xorBlock(chksum, chksum, tmp);                      // ... followed by the last (16 - ptLen) bytes of block-cipher output

		// Calculate and verify tag
		xorBlock(chksum, chksum, offset);
		tmp = encryptAes(chksum);
		for (int t = 0; t < TAG_LENGTH; t++) {
			if (tmp[t] != tag[0 + t]) {
				Arrays.fill(pt, (byte)0);
				pt = null;
				break;
			}
		}

		return pt;
	}

	//SecretKey skey;
	//SecretKeySpec skeySpec;
	byte[] key = new byte[16];

	/**
	 * Initialise the underlying AES cipher with a symmetric key.
	 *
	 * @param   cipherKey   the symmetric AES key.
	 * @param   keyBits     the key size in bits (128, 192, or 256).
	 *
	 * @throws  IllegalArgumentException    if the key is null or
	 *                      its length is not 128, 192, or 256.
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeySpecException 
	 */
	public final void init(byte[] k, int keyBits)
			throws IllegalArgumentException, NoSuchAlgorithmException, InvalidKeySpecException{
		if (k == null) {
			throw new IllegalArgumentException("Null key");
		}
		//if (keyBits != 128 && keyBits != 192 && keyBits != 256) {
		//throw new IllegalArgumentException("Invalid AES key size: " + keyBits + " bits.");
		//}
		// Initialize AES keys
		key = k;

		// Precompute L[i]-values (L[0] is synonym of L)
		byte[] L0 = L[0];
		Arrays.fill(L0, (byte)0);

		//aes.encrypt(L0, L0);
		L0 = encryptAes(L0);

		for (int i = 1; i < PRE_COMP_BLOCKS; i++) {
			// L[i] = L[i - 1] * x
			byte[] L_cur = L[i], L_pre = L[i - 1];
			for (int t = 0; t < BLOCK_SIZE - 1; t++) {
				L_cur[t] = (byte)((L_pre[t] << 1) | ((L_pre[t + 1] >>> 7) & 1));
			}
			L_cur[BLOCK_SIZE - 1] = (byte)((L_pre[BLOCK_SIZE - 1] << 1) ^ ((L_pre[0] & 0x80) != 0 ? 0x87 : 0));
		}

		// Precompute L_inv = L / x = L * x^{-1}
		for (int i = BLOCK_SIZE - 1; i > 0; i--) {
			L_inv[i] = (byte)(((L0[i] & 0xff) >>> 1) | ((L0[i - 1] & 1) << 7));
		}
		L_inv[0] = (byte)((L0[0] & 0xff) >>> 1);
		if ((L0[BLOCK_SIZE - 1] & 1) != 0) {
			L_inv[0] ^= 0x80;
			L_inv[BLOCK_SIZE - 1] ^= 0x43;
		}
	}

	protected static final void xorBlock(byte[] dst, int pos,
			byte[] src1, byte[] src2) {
		for (int i = 0; i < BLOCK_SIZE; i++) {
			dst[pos + i] = (byte)(src1[i] ^ src2[i]);
		}
	}

	protected static final void xorBlock(byte[] dst, byte[] src1,
			byte[] src2, int pos) {
		for (int i = 0; i < BLOCK_SIZE; i++) {
			dst[i] = (byte)(src1[i] ^ src2[pos + i]);
		}
	}

	protected static final void xorBlock(byte[] dst, byte[] src1, byte[] src2) {
		for (int i = 0; i < BLOCK_SIZE; i++) {
			dst[i] = (byte)(src1[i] ^ src2[i]);
		}
	}

	/**
	 * Count the number of trailing zeroes in integer i.
	 *
	 * @return  0 <= ntz <= 32.
	 */
	protected static final int ntz(int i) {
		/*
		 * N.B. this function is not as bad as it might seem:
		 * assuming the argument i is uniformly distributed,
		 * the probability that ntz(i) = k is 2^(-k-1),
		 * therefore the expected value of ntz(i) is < 1.
		 */
		int result = 0;
		while ((i & 1) == 0) {
			i >>>= 1;
		result++;
		}
		return result;
	}


	private byte[] encryptAes(byte[] src){
		byte[] result = null;
		try {
			SecretKeySpec skeySpec = new SecretKeySpec(key, "AES");
			Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");
			cipher.init(Cipher.ENCRYPT_MODE, skeySpec);
			result = cipher.doFinal(src);
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return result;
	}

	private byte[] decryptAes(byte[] src){
		byte[] result = null;
		try {
			//KeyGenerator kgen = KeyGenerator.getInstance("AES");
			//kgen.init(128); // 192 and 256 bits may not be available
			// Generate the secret key specs.
			//skey = kgen.generateKey();
			//byte[] raw = skey.getEncoded();
			SecretKeySpec skeySpec = new SecretKeySpec(key, "AES");
			Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");
			cipher.init(Cipher.DECRYPT_MODE, skeySpec);
			result = cipher.doFinal(src);
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return result;
	}

}
